{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fca5397d-d7dc-44d4-a888-3d8d04adc8db",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import datasets\n",
    "from dataclasses import dataclass\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acb2eb37-c0a3-48b7-8828-c96def0ac694",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\"\n",
    "model_path = \"./Qwen2.5-0.5B-SFT\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e37258e-f8c5-4f13-b7d0-5939913f2acb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_path,\n",
    "    dtype='auto',\n",
    "    device_map='auto'\n",
    ")\n",
    "\n",
    "ref_model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_path,\n",
    "    dtype='auto',\n",
    "    device_map='auto'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b02ce0d-3a84-4356-9e2b-33ed9224f1f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a42d66e-0010-4cc6-a940-f65b2df9297f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.generation_config.do_sample = True\n",
    "model.generation_config.eos_token_id = [151645, 151643]\n",
    "model.generation_config.pad_token_id = 151643\n",
    "model.generation_config.temperature = 0.7\n",
    "model.generation_config.top_p = 0.8\n",
    "model.generation_config.top_k = 20\n",
    "model.generation_config.repetition_penalty = 1.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c57c11d-3546-4288-9e91-94c3b0585779",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class DPOConfig:\n",
    "    max_length: int = 1700  # 根据自身具备的算力条件进行自适应更改\n",
    "    batch_size: int = 2\n",
    "    gradient_accumulation_steps: int = 8\n",
    "    beta: float = 0.5\n",
    "    log_iter: int = 200\n",
    "    max_lr: float = 1e-6\n",
    "    min_lr: float = 1e-7\n",
    "    warmup_steps: int = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f1ab3d5-b852-4268-8a32-74fc90d1cb98",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
